from rapidocr import RapidOCR
import requests
from PIL import Image
import numpy as np
import io
import os
import re
import fitz  # PyMuPDF

def download_file(url):
    """下载文件并返回字节内容"""
    response = requests.get(url, timeout=30)
    response.raise_for_status()
    return response.content

def download_image(url):
    """下载图片并返回PIL Image对象"""
    response = requests.get(url, timeout=10)
    response.raise_for_status()
    return Image.open(io.BytesIO(response.content))

def extract_pdf_pages_as_images(pdf_content, page_range="1-1"):
    """
    从PDF内容中提取指定页面作为图片
    
    Args:
        pdf_content (bytes): PDF文件的字节内容
        page_range (str): 页面范围，格式为"start-end"或"page"
    
    Returns:
        list: PIL Image对象列表
    """
    # 解析页面范围
    if "-" in page_range:
        start, end = map(int, page_range.split("-"))
    else:
        start = end = int(page_range)
    
    # 确保起始页码至少为1
    start = max(1, start)
    end = max(start, end)
    
    images = []
    try:
        # 打开PDF文档
        pdf_document = fitz.open(stream=pdf_content, filetype="pdf")
        total_pages = pdf_document.page_count
        
        print(f"PDF总页数: {total_pages}")
        print(f"提取页面范围: {start}-{end}")
        
        # 确保页面范围不超过PDF总页数
        end = min(end, total_pages)
        
        # 提取指定页面
        for page_num in range(start-1, end):  # fitz使用0基索引
            page = pdf_document[page_num]
            # 将页面转换为图片
            pix = page.get_pixmap(dpi=200)  # 设置DPI以获得更好的图像质量
            img_data = pix.tobytes("ppm")
            img = Image.open(io.BytesIO(img_data))
            images.append(img)
            print(f"已提取第 {page_num+1} 页")
        
        pdf_document.close()
    except Exception as e:
        print(f"PDF页面提取失败: {e}")
    
    return images

def extract_text_from_ocr_results(ocr_results):
    """
    从OCR结果中提取文本内容
    
    Args:
        ocr_results (list): OCR识别结果列表
    
    Returns:
        str: 提取的所有文本内容
    """
    all_text = ""
    
    for page_result in ocr_results:
        page_num = page_result['page']
        result = page_result['result']
        
        # 根据测试结果，RapidOCR返回的是RapidOCROutput对象
        try:
            # 检查是否有txts属性且不为None
            if hasattr(result, 'txts') and result.txts is not None:
                # 从txts属性提取文本
                if isinstance(result.txts, list):
                    # 过滤掉空字符串和None值
                    valid_texts = [text for text in result.txts if text]
                    if valid_texts:
                        all_text += " ".join(valid_texts) + " "
                        print(f"从第 {page_num} 页提取到 {len(valid_texts)} 个文本片段")
                else:
                    # 如果txts不是列表，直接转换为字符串
                    text = str(result.txts)
                    if text and text.strip():  # 只有非空文本才添加
                        all_text += text + " "
                        print(f"从第 {page_num} 页提取到文本: {text}")
            else:
                print(f"第 {page_num} 页没有识别到文本")
        except Exception as e:
            print(f"处理第 {page_num} 页文本时出错: {e}")
    
    return all_text.strip()

def ocr_from_image_array(img_array):
    """对numpy数组图像进行OCR识别"""
    try:
        # 初始化OCR引擎
        engine = RapidOCR()
        
        # 执行OCR
        result = engine(img_array)
        
        return result
    except Exception as e:
        print(f"OCR处理失败: {e}")
        return None

def ocr_single_page(image, page_index=0):
    """对单个页面图像进行OCR识别"""
    try:
        print(f"正在进行第 {page_index+1} 页OCR识别...")
        
        # 转换为numpy数组
        img_array = np.array(image)
        
        # 执行OCR
        result = ocr_from_image_array(img_array)
        
        if result:
            print(f"第 {page_index+1} 页OCR识别完成")
            # 根据测试结果，检查txts属性
            if hasattr(result, 'txts') and result.txts is not None:
                # 过滤掉空字符串和None值
                if isinstance(result.txts, list):
                    valid_texts = [text for text in result.txts if text]
                    if valid_texts:
                        print(f"  识别到 {len(valid_texts)} 个文本片段")
                        for i, text in enumerate(valid_texts[:5]):  # 只显示前5个
                            print(f"    文本 {i+1}: {text}")
                        if len(valid_texts) > 5:
                            print(f"    ... 还有 {len(valid_texts) - 5} 个文本片段")
                    else:
                        print(f"  未识别到有效文本")
                else:
                    text = str(result.txts)
                    if text and text.strip():
                        print(f"  识别到文本: {text}")
                    else:
                        print(f"  未识别到有效文本")
            else:
                print(f"  未识别到文本")
            return result
        else:
            print(f"第 {page_index+1} 页OCR识别失败或未返回结果")
        
        return result
    except Exception as e:
        print(f"第 {page_index+1} 页OCR处理失败: {e}")
        return None

def process_pdf_and_extract_text(pdf_url, page_range="1-1", find_stop=True):
    """
    下载PDF文件，提取指定页面并进行OCR识别
    
    Args:
        pdf_url (str): PDF文件URL
        page_range (str): 页面范围，默认为"1-1"
        find_stop (bool): 是否在找到链接后停止识别后续页面，默认为True
    
    Returns:
        list: 所有页面的OCR识别结果
    """
    print(f"正在处理PDF文件: {pdf_url}")
    print(f"页面范围: {page_range}")
    print(f"找到链接后停止识别: {find_stop}")
    
    try:
        # 下载PDF文件
        print("正在下载PDF文件...")
        pdf_content = download_file(pdf_url)
        
        # 提取指定页面作为图片
        print("正在提取PDF页面...")
        images = extract_pdf_pages_as_images(pdf_content, page_range)
        
        if not images:
            print("未能提取到任何页面")
            return []
        
        # 对每一页进行OCR识别
        all_results = []
        
        for i, image in enumerate(images):
            # 如果启用了find_stop并且已经处理了至少一页，检查是否需要停止
            if find_stop and all_results:
                # 检查之前页面是否已经找到了链接
                prev_result = all_results[-1]['result']
                if hasattr(prev_result, 'txts') and prev_result.txts is not None:
                    # 从之前的结果中提取文本
                    prev_text = ""
                    if isinstance(prev_result.txts, list):
                        valid_texts = [text for text in prev_result.txts if text]
                        prev_text = " ".join(valid_texts)
                    else:
                        prev_text = str(prev_result.txts)
                    
                    # 检查之前页面是否包含链接
                    extracted_links = extract_links_with_keywords(prev_text, ["https", ".com", ".io"])
                    if extracted_links:
                        print(f"在第 {all_results[-1]['page']} 页找到链接，停止识别后续页面")
                        break
            
            result = ocr_single_page(image, i)
            if result:
                all_results.append({
                    'page': i + 1,
                    'result': result
                })
                
                # 如果启用了find_stop并且这是第一页，立即检查是否找到链接
                if find_stop and len(all_results) == 1:
                    if hasattr(result, 'txts') and result.txts is not None:
                        # 从当前结果中提取文本
                        current_text = ""
                        if isinstance(result.txts, list):
                            valid_texts = [text for text in result.txts if text]
                            current_text = " ".join(valid_texts)
                        else:
                            current_text = str(result.txts)
                        
                        # 检查当前页面是否包含链接
                        extracted_links = extract_links_with_keywords(current_text, ["https", ".com", ".io"])
                        if extracted_links:
                            print(f"在第 {i+1} 页找到链接，停止识别后续页面")
                            break
        
        return all_results
        
    except Exception as e:
        print(f"PDF处理失败: {e}")
        return []

def extract_links_with_keywords(text, contain_string=["https", ".com", ".io"]):
    """
    从文本中提取包含指定关键字的链接
    
    Args:
        text (str): 要搜索的文本
        contain_string (list): 包含的关键字列表，默认为["https", ".com", ".io"]
    
    Returns:
        list: 包含指定关键字的链接列表
    """
    print(f"正在从文本中提取链接，查找包含 {contain_string} 的链接...")
    
    # 如果text是元组或其他非字符串类型，先转换为字符串
    if not isinstance(text, str):
        text = str(text)
    
    print(f"处理的文本长度: {len(text)} 字符")
    
    # 使用多种正则表达式查找URL
    # 1. 匹配http/https链接（完整URL）
    https_pattern = r'https://[^\s<>"]+'
    http_pattern = r'http://[^\s<>"]+'
    
    # 2. 匹配www开头的链接
    www_pattern = r'www\.[^\s<>"]+'
    
    # 3. 匹配域名格式（如 robot-colosseum.github.io）
    domain_pattern = r'[a-zA-Z0-9\-\.]+\.(?:io|com|org|edu|gov|net)(?:/[^\s<>"]*)?'
    
    # 查找所有匹配的URL
    https_urls = re.findall(https_pattern, text)
    http_urls = re.findall(http_pattern, text)
    www_urls = re.findall(www_pattern, text)
    domain_urls = re.findall(domain_pattern, text)
    
    # 合并所有潜在URL并去重
    potential_urls = list(set(https_urls + http_urls + www_urls + domain_urls))
    
    print(f"找到 {len(potential_urls)} 个潜在URL")
    
    # 过滤包含指定关键字的URL
    filtered_urls = []
    
    for url_str in potential_urls:
        print(f"检查URL: {url_str}")
        
        # 检查是否包含任何指定的关键字
        contains_keyword = any(keyword in url_str for keyword in contain_string)
        
        if contains_keyword:
            print(f"  URL包含关键字: {url_str}")
            
            # 标准化URL格式
            if url_str.startswith("www."):
                url_str = "https://" + url_str
            elif not url_str.startswith(("http://", "https://")):
                # 对于域名格式，添加https://前缀
                if "." in url_str:
                    url_str = "https://" + url_str
            
            # 清理URL末尾的标点符号
            # 链接的最后一位可以是/，英文与数字组合，但不能是.或'等符号
            while url_str and url_str[-1] in '.,;!?\'"()[]{}':
                url_str = url_str[:-1]
                
            # 确保URL是有效的格式
            if url_str.startswith(("http://", "https://")) and "." in url_str:
                filtered_urls.append(url_str)
        else:
            print(f"  URL不包含指定关键字: {url_str}")
    
    # 去重并保持顺序
    unique_urls = []
    for url in filtered_urls:
        # 清理URL末尾的标点符号
        # 链接的最后一位可以是/，英文与数字组合，但不能是.或'等符号
        while url and url[-1] in '.,;!?\'"()[]{}':
            url = url[:-1]
        
        # 确保是有效的URL格式
        if url.startswith(("http://", "https://")) and "." in url:
            # 检查是否已存在（避免重复）
            if url not in unique_urls:
                unique_urls.append(url)
    
    print(f"最终提取到 {len(unique_urls)} 个符合条件的链接")
    return unique_urls

def ocr_and_extract_links(file_url, contain_string=["https", ".com", ".io"], page_range="1-2", save_vis=True, find_stop=True):
    """
    下载文件URL（支持PDF和图片），使用OCR识别内容，并提取包含指定关键字的链接
    
    Args:
        file_url (str): 文件URL地址（支持PDF和图片）
        contain_string (list): 包含的关键字列表，默认为["https", ".com", ".io"]
        page_range (str): 页面范围，默认为"1-2"
        save_vis (bool): 是否保存可视化结果，默认为True
        find_stop (bool): 是否在找到链接后停止识别后续页面，默认为True
    
    Returns:
        dict: 包含OCR结果和提取链接的信息
    """
    print(f"正在处理文件URL: {file_url}")
    print(f"页面范围: {page_range}")
    print(f"找到链接后停止识别: {find_stop}")
    
    # 检查是否为PDF文件
    if file_url.lower().endswith('.pdf'):
        # 处理PDF文件
        ocr_results = process_pdf_and_extract_text(file_url, page_range, find_stop)
    else:
        # 处理图片文件
        print("检测到图片文件，正在处理...")
        try:
            # 下载图片
            img = download_image(file_url)
            # 进行OCR识别
            result = ocr_single_page(img, 0)
            ocr_results = [{'page': 1, 'result': result}] if result else []
        except Exception as e:
            print(f"图片处理失败: {e}")
            ocr_results = []
    
    if not ocr_results:
        print("OCR识别失败或未识别到任何文本")
        return {
            "url": file_url,
            "ocr_success": False,
            "extracted_links": [],
            "page_range": page_range
        }
    
    # 从OCR结果中提取文本
    all_text = extract_text_from_ocr_results(ocr_results)
    
    print(f"OCR识别到的文本内容长度: {len(all_text)} 字符")
    # 显示前500个字符用于调试
    print(f"OCR识别到的文本内容 (前500字符): {all_text[:500]}...")
    
    # 从OCR文本中提取链接
    extracted_links = extract_links_with_keywords(all_text, contain_string)
    
    result_info = {
        "url": file_url,
        "ocr_success": True,
        "extracted_links": extracted_links,
        "total_links_found": len(extracted_links),
        "page_range": page_range,
        "contains_keywords": contain_string,
        "ocr_text": all_text,  # 添加完整文本用于调试
        "pages_processed": len(ocr_results)  # 添加处理的页面数
    }
    
    # 打印提取到的链接
    if extracted_links:
        print(f"\n从OCR文本中提取到 {len(extracted_links)} 个链接:")
        for i, link in enumerate(extracted_links, 1):
            print(f"  {i}. {link}")
    else:
        print("\n未从OCR文本中提取到符合条件的链接")
        # 额外检查：在完整文本中搜索包含关键字的文本
        print("\n额外检查：在完整文本中搜索包含关键字的文本...")
        for keyword in contain_string:
            if keyword in all_text:
                print(f"  发现包含 '{keyword}' 的文本")
                # 显示包含关键字的前后各50个字符
                index = all_text.find(keyword)
                start = max(0, index - 50)
                end = min(len(all_text), index + 100)
                print(f"    上下文: ...{all_text[start:end]}...")
    
    return result_info

# 使用示例
if __name__ == "__main__":
    # 测试PDF URL处理
    pdf_url = "https://www.roboticsproceedings.org/rss20/p133.pdf"
    
    print("测试PDF文件处理和链接提取:")
    print("="*50)
    
    result = ocr_and_extract_links(
        file_url=pdf_url,
        page_range="1-2",  # 处理前2页
        contain_string=["https", ".com", ".io"],
        find_stop=True  # 找到链接后停止识别
    )
    
    print("\n" + "="*50)
    print("处理结果摘要:")
    print("="*50)
    print(f"文件URL: {result['url']}")
    print(f"OCR识别成功: {result['ocr_success']}")
    print(f"提取到链接数: {result['total_links_found']}")
    print(f"页面范围: {result['page_range']}")
    print(f"处理页面数: {result.get('pages_processed', 'N/A')}")
    
    if result['extracted_links']:
        print("提取到的链接:")
        for i, link in enumerate(result['extracted_links'], 1):
            print(f"  {i}. {link}")
    else:
        print("未提取到符合条件的链接")
